# -*- coding: utf-8 -*
# cython: profile=True
from numpy import array,dot,cross,ndarray,isnan,argmin,argmax
from numpy.random import random
from random import sample
from math import sqrt,sin,cos,exp
import cython
if not cython.compiled:
    from math import *
from v3dsolver.vector import Vec
import pdb,time

def eval_legendre(coeffs,x):
    """
    Evaluates legendre polynomial specified by coeffs at x.
    
    coeffs: array of coefficients in increasing order
    x: point at which legendre is to be evaluated.
    """
    result = 0.
    for i in range(len(coeffs)):
        coeff = coeffs[i]
        result += coeff*(x**(i))
    return result

class FreeStream:
    """
    Generic free stream object. A convenient hook to handle various types of 
    free stream conditions.
    """
    def __init__(self,freestream_val):
        """
        Initialization of free stream object. This implementation assumes a 
        constant free stream condition given by freestream_val.
        """
        if isinstance(freestream_val,Vec):
            self.freestream_val = freestream_val
        else:
            self.freestream_val = Vec(*freestream_val)
        
    def get_vel(self,fluid_elem):
        """
        Returns the velocity on fluid_elem due to free stream. Returns constant 
        freestream_val for current case.
        """
        return self.freestream_val

class FluidElem:
    """
    Class representing a generic fluid element. A fluid element is an entity of 
    fluid which has position and velocity.
    """
    def __init__(self,pos,vel,cause_or_effect):
        """
        Initialization of fluid element using its pos and vel. The variable 
        cause_or_effect is:
        
        0: for pure cause fluid elements (e.g. vortex panel, boundary blob, 
           pseudo particle)
        1: for pure effect fluid elements (e.g. stretching stensil point, 
           control point)
        2: for both cause and effect fluid elements (e.g. any vortex blob)
        """
        if isinstance(pos,Vec):
            self.pos = pos
        else:
            self.pos = Vec(*pos)
        if isinstance(vel,Vec):
            self.vel = vel
        else:
            self.vel = Vec(*vel)
        self.cause_or_effect = cause_or_effect
        
    def move(self,time_interval):
        """
        Move the fluid element for time_interval with its current velocity vel
        """
        self.pos += self.vel*time_interval
        
    def is_cause(self):
        """
        Returns True if fluid element is pure cause or both cause and effect. 
        Otherwise, returns False.
        """
        if self.cause_or_effect==0 or self.cause_or_effect==2:
            return True
        else:
            return False
        
    def is_effect(self):
        """
        Returns True if fluid element is pure effect or both cause and effect. 
        Otherwise, returns False.
        """
        if self.cause_or_effect==1 or self.cause_or_effect==2:
            return True
        else:
            return False

    def __reduce__(self):
        """
        Implemented to facililate pickling of extension types.
        """
        d = {}
        d['pos'] = self.pos
        d['vel'] = self.vel
        d['cause_or_effect'] = self.cause_or_effect
        return (FluidElem,(self.pos,self.vel,self.cause_or_effect),d)

    def __setstate__(self,d):
        """
        Load the FluidElem object from the saved dictionary.
        """
        self.pos = d['pos']
        self.vel = d['vel']
        self.cause_or_effect = d['cause_or_effect']

class StretchStensilPoint(FluidElem):
    """
    Class representing a vertex of the stensil surrounding a vortex blob used 
    to calculate stretch parameters.
    """
    def __init__(self,parent_particle,direction,dx):
        """
        Initialization of StretchStensilPoint object belonging to 
        parent_particle. This object is used to calculate velocity at a point 
        dx ahead of parent_particle along direction given by direction.
        
        The parameter direction can have the following values:
        0: along x-direction
        1: along y-direction
        2: along z-direction
        """
        self.parent_particle = parent_particle
        self.direction = direction
        pos = parent_particle.pos + dx*self.direction
        vel = Vec()# initializing stretch stensil pt with zero vel
        FluidElem.__init__(self,pos,vel,cause_or_effect=1)
        self.parentid = [] # by default, doesn't have a parent
    
    def __reduce__(self):
        """
        Implemented to facililate pickling of extension types.
        """
        d = {}
        d1 = FluidElem.__reduce__(self)[2]
        d.update(d1)
        d['parent_particle'] = self.parent_particle
        d['direction'] = self.direction
        d['parentid'] = self.parentid
        return (StretchStensilPoint,(self.parent_particle,self.direction,None),d)

    def __setstate__(self,d):
        """
        Load the FluidElem object from the saved dictionary.
        """
        FluidElem.__setstate__(self,d)
        self.parent_particle = d['parent_particle']
        self.direction = d['direction']
        self.parentid = d['parentid']


class VortexPanel(FluidElem):

    """
    Class representing a triangular vortex panel.
    """
    def __init__(self,vertices,rad,vel,delta):
        """
        Initialization of a panel made up of three points contained in the list 
        'vertices'. The order of points in vertices should be such that the 
        right hand thumb rule gives outward normal to the body.
        
        rad: core radius of boundary blobs contained by the vortex panel. 
        vel: velocity of panel (defaulted to zero for non-moving bodies)
        delta: distance of boundary blob beneath the surface
        """
        self.vertices = [Vec(vertex[0],vertex[1],vertex[2]) for vertex in vertices]
        normal = (self.vertices[1]-self.vertices[0]).cross(
                          self.vertices[2]-self.vertices[1])# outward normal
        self.unitnormal = normal/normal.length() # unit normal
        self.centroid = Vec()
        for vertex in self.vertices:
            self.centroid += vertex
        self.centroid /= float(len(self.vertices))
        FluidElem.__init__(self,pos=self.centroid,vel=vel,cause_or_effect=0)
        self.rad = rad
        self.vor = Vec() # vorticity unknown yet
        self.delta = delta
        self.ns_controlpt, self.np_controlpts = self.get_ctrlpts()
        self.cause_pos = self.ns_controlpt.pos # pos of cause elems
        self.bblob = BoundaryBlob(parent_panel=self)
        self.rblob = ReleaseBlob(parent_panel=self)
        self.cause_elems = [self.bblob,self.rblob]
        self.basis = self.get_basis()
    
    def get_ctrlpts(self):
        """
        Returns no-slip controlpt and no-penetration controlpts. Control 
        points are placed halfway across the lines joining the centroid 
        and the vertices. No-slip controlpt is placed opposite to the 
        smallest side and the remaining two positions are filled with 
        no-penetration controlpts.
        """
        ##
        
        ns_ctrlpt = ControlPt(pos=self.centroid,parent_panel=self)
        
        np_ctrlpts = []
        for vertex in self.vertices:
            np_pos = self.centroid + 0.5*(vertex - self.centroid)
            np_ctrlpts += [ControlPt(pos=np_pos,parent_panel=self)]
        
        ##
        """
        sides = {0:(self.vertices[2]-self.vertices[1]).length(),
                 1:(self.vertices[0]-self.vertices[2]).length(),
                 2:(self.vertices[1]-self.vertices[0]).length()}
        args_sorted = sorted(sides,key=sides.__getitem__)
        
        i_ns = args_sorted[0]
        i_np1 = args_sorted[1]
        i_np2 = args_sorted[2]
        
        ns_pos = self.centroid  + 0.5*(self.vertices[i_ns] - self.centroid)
        ns_ctrlpt = ControlPt(pos=ns_pos,parent_panel=self)
        
        np1_pos = self.centroid + 0.5*(self.vertices[i_np1] - self.centroid)
        np2_pos = self.centroid + 0.5*(self.vertices[i_np2] - self.centroid)
        np_ctrlpts = [ControlPt(pos=np1_pos,parent_panel=self)]
        np_ctrlpts += [ControlPt(pos=np2_pos,parent_panel=self)]
        
        ##
                 
        # obtain no-slip controlpt
        i_ctrlpts = range(3)
        i_ns = argmin(sides)
        ns_pos = self.centroid
        ns_ctrlpt = ControlPt(pos=ns_pos,parent_panel=self)
        
        # obtain no penetration controlpt
        i_ctrlpts.remove(i_ns)
        i_np1,i_np2 = sample(range(3),2)
        np1_pos = self.centroid + 0.5*(self.vertices[i_np1] - self.centroid)
        np2_pos = self.centroid + 0.5*(self.vertices[i_np2] - self.centroid)
        np_ctrlpts = [ControlPt(pos=np1_pos,parent_panel=self)]
        np_ctrlpts += [ControlPt(pos=np2_pos,parent_panel=self)]
        """
        
        return ns_ctrlpt,np_ctrlpts
    
    def get_basis(self):
        """
        Returns two orthogonal unit vectors along the panel plane.
        """
        nx,ny,nz = self.unitnormal.x,self.unitnormal.y,self.unitnormal.z
        vec1 = Vec(-nz,-nz,nx+ny) # perpendicular to normal
        vec1 = vec1/vec1.length()
        vec2 = vec1.cross(self.unitnormal) # perpendicular to vec1 & normal
        vec2 = vec2/vec2.length()
        return [vec1,vec2]
        
    def __reduce__(self):
        """
        Implemented to facililate pickling of extension types.
        """
        d = {}
        d1 = FluidElem.__reduce__(self)[2]
        d.update(d1)
        d['vertices'] = self.vertices
        d['unitnormal'] = self.unitnormal
        d['centroid'] = self.centroid
        d['rad'] = self.rad
        d['vor'] = self.vor
        d['delta'] = self.delta
        d['cause_pos'] = self.cause_pos
        d['bblob'] = self.bblob
        d['rblob'] = self.rblob
        d['cause_elems'] = self.cause_elems
        d['np_controlpts'] = self.np_controlpts
        d['ns_controlpt'] = self.ns_controlpt
        d['basis'] = self.basis
        return (VortexPanel,(self.vertices,self.rad,self.vel,self.delta),d)

    def __setstate__(self,d):
        """
        Load the FluidElem object from the saved dictionary.
        """
        FluidElem.__setstate__(self,d)
        self.vertices = d['vertices']
        self.unitnormal = d['unitnormal']
        self.centroid = d['centroid']
        self.rad = d['rad']
        self.vor = d['vor']
        self.delta = d['delta']
        self.cause_pos = d['cause_pos']
        self.bblob = d['bblob']
        self.rblob = d['rblob']
        self.cause_elems = d['cause_elems']
        self.np_controlpts = d['np_controlpts']
        self.ns_controlpt = d['ns_controlpt']
        self.basis = d['basis']

class VortexBlob(FluidElem):
    """
    Class representing a generic vortex blob object. Specific blob types: point vortex, saffman, krasny 
    etc. will be derived from this class.
    """
    def __init__(self,pos,vor,vel,rad):
        """
        Initialization of vortex blob using its pos, vel, vor (vorticity) and rad (core radius)
        """
        FluidElem.__init__(self,pos,vel,cause_or_effect=2)
        self.rad = rad
        if isinstance(vor,Vec):
            self.vor = vor
        else:
            self.vor = Vec(*vor)
        self.parentid = [] # by default, doesn't have a parent
        self.stretchpts = [] # by default, doesn't have any stretch stensil

    def __reduce__(self):
        """
        Implemented to facililate pickling of extension types.
        """
        d = {}
        d1 = FluidElem.__reduce__(self)[2]
        d.update(d1)
        d['rad'] = self.rad
        d['vor'] = self.vor
        d['parentid'] = self.parentid
        d['stretchpts'] = self.stretchpts
        return (VortexBlob,(self.pos,self.vor,self.vel,self.rad),d)

    def __setstate__(self,d):
        """
        Load the FluidElem object from the saved dictionary.
        """
        FluidElem.__setstate__(self,d)
        self.rad = d['rad']
        self.vor = d['vor']
        self.parentid = d['parentid']
        self.stretchpts = d['stretchpts']
    
class Saffman(VortexBlob):
    """
    Class representing Saffman vortex blob.
    """
    def __init__(self,pos,vor,vel,rad):
        VortexBlob.__init__(self,pos,vor,vel,rad)
        
    def __reduce__(self):
        """
        Implemented to facililate pickling of extension types.
        """
        d = VortexBlob.__reduce__(self)[2]
        return (Saffman,(self.pos,self.vor,self.vel,self.rad),d)

    def get_vel(self,fluid_elem):
        """ 
        Returns velocity induced by this Saffman vortex blob on given 
        fluid_elem. If 'fluid_elem' is of type FluidElem,velocity at its 
        position is returned, else if 'fluid_elem' is of type length-3 
        array (pos), velocity at this pos is returned. 
        """ 
        # skipping velocity calculation on a stretch point due to it's own parent
        vel = Vec()
        if fluid_elem in self.stretchpts:
            return vel
        rvec = fluid_elem.pos - self.pos
        r = rvec.length()
        sigma = self.rad
        if r==0.0: # if vel on self is evaluated
            K_delta = Vec()
        else:
            pi = 3.1415926535897931
            f = tanh((r/sigma)**3)
            f = f + 1.5*((r/sigma)**3)*(1-f**2) # 4th order kernel
            K_delta = rvec.mul(-f/(4.*pi*r**3))
        vel = K_delta.cross(self.vor)
        return vel

class PseudoParticle(Saffman):
    """
    Class representing a pseudo particle generated as a result of Makino's 
    scheme on far-off cluster of particles.
    """
    def __init__(self,pos,parent_cell,rad,P_l,particles=[]):
        """
        Initialization of the PseudoParticle object. Parameters pos, vor 
        represent position and vorticity of pseudo particle, parent_cell 
        represents the cell object containing the pseudo particle. 
        
        P_l: a list of legendre polynomial functions
        particles: list of cause elems
                   if [], vor of pseudo particle would be calculated based on 
                   parent_cell.cause_elems; otherwise, it would be based on 
                   these particles
        rad: radius of PseudoParticle blob
        """
        FluidElem.__init__(self,pos,vel=Vec(),cause_or_effect=0)
        self.P_l = P_l
        if particles==[]:
            particles = parent_cell.cause_elems
        self.elems = particles
        self.vor = self.get_vor(parent_cell)
        Saffman.__init__(self,self.pos,self.vor,self.vel,rad)
    
    def __reduce__(self):
        """
        Implemented to facililate pickling of extension types.
        """
        d = {}
        d1 = Saffman.__reduce__(self)[2]
        d.update(d1)
        d['degree_lim'] = self.degree_lim
        d['elems'] = self.elems
        return (PseudoParticle,(self.pos,None,self.rad,self.degree_lim,None),d)

    def __setstate__(self,d):
        """
        Load the FluidElem object from the saved dictionary.
        """
        Saffman.__setstate__(self,d)
        self.degree_lim = d['degree_lim']
        self.elems = d['elems']

    def get_vor(self,cell):
        """
        Returns vorticity of the pseudo particle based on Makino's method.
        """
        vor = Vec()
        center2pseudo = self.pos - cell.center
        cell_center = cell.center
        cell_radius = cell.radius
        for elem in self.elems:
            center2elem = elem.pos - cell_center
            center2elem_length = center2elem.length()
            factor1 = (center2elem_length/cell_radius)
            elem_vor = elem.vor
            if center2elem.length()>=1e-7:          # skipping the case when elem is AT the cell center       
                cos_angle = (center2elem.dot(center2pseudo)/
                            (center2elem.length()*center2pseudo.length()))
                for l in range(len(self.P_l)):
                    legendre = eval_legendre(self.P_l[l],cos_angle)
                    factor = (2*l+1)*(factor1**l)*legendre
                    vor += elem_vor*factor
        vor = vor/float(cell.npseudo)
        return vor

class BoundaryBlob(Saffman):
    """
    Class representing a boundary blob object associated with a given vortex panel.
    """
    def __init__(self,parent_panel):
        """
        Initialization of a boundary blob object having pos, vor(vorticity), 
        rad(core radius) and parent_panel(associated vortex panel). Velocity 
        of boundary blob is same as that of parent panel. This is a pure cause 
        element.
        """
        self.parent_panel = parent_panel
        pos = parent_panel.cause_pos - parent_panel.delta*(parent_panel.unitnormal)
        Saffman.__init__(self,pos=pos,vor=Vec(),
                         vel=parent_panel.vel,rad=parent_panel.rad)
        self.cause_or_effect = 0

    def __reduce__(self):
        """
        Implemented to facililate pickling of extension types.
        """
        d = {}
        d1 = Saffman.__reduce__(self)[2]
        d.update(d1)
        d['parent_panel'] = self.parent_panel
        return (BoundaryBlob,(self.parent_panel,),d)

    def __setstate__(self,d):
        """
        Load the FluidElem object from the saved dictionary.
        """
        Saffman.__setstate__(self,d)
        self.parent_panel = d['parent_panel']
        
class ReleaseBlob(Saffman):
    """
    Class representing blobs which are just release from the boundary to 
    satisfy no-slip. 
    """
    def __init__(self,parent_panel):
        """
        Initialization of ReleaseBlob object. The position of this object is 
        delta distance along outward unit normal above corresponding copy_blob.
        Vorticity is such that normal component is reversed and tangential
        component is retained.
        """
        self.parent_panel = parent_panel
        self.copy_blob = parent_panel.bblob
        pos = parent_panel.cause_pos + parent_panel.delta*(parent_panel.unitnormal)
        Saffman.__init__(self,pos=pos,vor=Vec(),
                         vel=parent_panel.vel,rad=parent_panel.rad)
        self.cause_or_effect = 0

    def __reduce__(self):
        """
        Implemented to facililate pickling of extension types.
        """
        d = {}
        d1 = Saffman.__reduce__(self)[2]
        d.update(d1)
        d['copy_blob'] = self.copy_blob
        d['parent_panel'] = self.parent_panel
        return (ReleaseBlob,(self.parent_panel,),d)

    def __setstate__(self,d):
        """
        Load the FluidElem object from the saved dictionary.
        """
        Saffman.__setstate__(self,d)
        self.copy_blob = d['copy_blob']
        self.parent_panel = d['parent_panel']

class ControlPt(Saffman):
    """
    Class representing a boundary blob object associated with a given vortex panel.
    """
    def __init__(self,pos,parent_panel):
        """
        Initialization of a boundary blob object having pos, vor(vorticity), 
        rad(core radius) and parent_panel(associated vortex panel). Velocity 
        of boundary blob is same as that of parent panel. This is a pure cause 
        element.
        """
        self.parent_panel = parent_panel
        Saffman.__init__(self,pos=pos,vor=Vec(),
                         vel=parent_panel.vel,rad=parent_panel.rad)
        self.cause_or_effect = 1

    def __reduce__(self):
        """
        Implemented to facililate pickling of extension types.
        """
        d = {}
        d1 = Saffman.__reduce__(self)[2]
        d.update(d1)
        d['parent_panel'] = self.parent_panel
        return (ControlPt,(self.pos,self.parent_panel),d)

    def __setstate__(self,d):
        """
        Load the FluidElem object from the saved dictionary.
        """
        Saffman.__setstate__(self,d)
        self.parent_panel = d['parent_panel']
